import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression #现行回归模型
from sklearn.preprocessing import PolynomialFeatures #构建多项式特征
from sklearn.model_selection import train_test_split #花粉训练集和测试集
from sklearn.metrics import mean_squared_error #均方误差损失函数

"""
1.生成数据
2.划分训练接和测试集
3.定义模型（线性回归模型）
4.训练模型
5.预算结果，计算误差
"""
plt.rcParams['font.sans-serif']=['KaiTi']
plt.rcParams['axes.unicode_minus']=False

#1.生成数据
X = np.linspace(-3,3,300).reshape(-1,1)
Y = np.sin(X) + np.random.uniform(low=-0.5,high=0.5,size=300).reshape(-1,1)
print(X.shape)
print(Y.shape)
#画出散点图（三个子图）
fig,ax = plt.subplots(1,3,figsize = (15,4))
ax[0].scatter(X,Y,c='y')
ax[1].scatter(X,Y,c='y')
ax[2].scatter(X,Y,c='y')


#2.划分训练集和测试集
trainX,testX,trainY,testY=train_test_split(X,Y,test_size=0.2,random_state=42)

#3.定义模型（现行回归模型）
model1 = LinearRegression()
model2 = LinearRegression()
model3 = LinearRegression()
#一、欠拟合（直线）
x_train1 = trainX
x_test1 = testX
#恰好拟合（5次多项式）
poly5 = PolynomialFeatures(degree=5)
x_train2 = poly5.fit_transform(trainX)
x_test2 = poly5.fit_transform(testX)
#过拟合（20次多项式）
poly20 = PolynomialFeatures(degree=20)
x_train3 = poly20.fit_transform(trainX)
x_test3 = poly20.fit_transform(testX)
#4.训练模型
model1.fit(x_train1,trainY)
model2.fit(x_train2,trainY)
model3.fit(x_train3,trainY)

#打印查看模型参数
print(model1.coef_)
print(model1.intercept_)

#5.预测结果，计算误差
y_pred1 = model1.predict(x_test1)
test_loss1 = mean_squared_error(testY,y_pred1)
train_loss1 = mean_squared_error(trainY,model1.predict(x_train1))

y_pred2 = model2.predict(x_test2)
test_loss2 = mean_squared_error(testY,y_pred2)
train_loss2 = mean_squared_error(trainY,model2.predict(x_train2))

y_pred3 = model3.predict(x_test3)
test_loss3 = mean_squared_error(testY,y_pred3)
train_loss3 = mean_squared_error(trainY,model3.predict(x_train3))

#画出拟合曲线并写出训练误差和测试误差
ax[0].plot(X,model1.predict(X),c='r')
ax[0].text(-3,1,f"测试误差：{test_loss1:.4f}")
ax[0].text(-3,1.3,f"训练误差{train_loss1:.4f}")

ax[1].plot(X,model2.predict(poly5.fit_transform(X)),c='r')
ax[1].text(-3,1,f"测试误差：{test_loss2:.4f}")
ax[1].text(-3,1.3,f"训练误差{train_loss2:.4f}")

ax[2].plot(X,model3.predict(poly20.fit_transform(X)),c='r')
ax[2].text(-3,1,f"测试误差：{test_loss3:.4f}")
ax[2].text(-3,1.3,f"训练误差{train_loss3:.4f}")
plt.show()



